import tensorrt as trt
import pycuda.autoinit  # This is needed for initializing CUDA driver
import pycuda.driver as cuda
import cv2
import numpy as np
import ctypes
from threading import Thread,Lock,Condition
from camera import Camera
import subprocess
import logging
import os
from utils.cry import callback

class TrtSSD(Thread):
    """TrtSSD class encapsulates things needed to run TRT SSD."""

    def __load_plugins(self):
        ctypes.CDLL("./utils/libflattenconcat.so")
        trt.init_libnvinfer_plugins(self.__trt_logger, '')
    def __load_engine(self):
        TRTbin = 'networks/TRT_%s.bin' % self.model
        with trt.Runtime(self.__trt_logger) as runtime:
            return  callback(runtime.deserialize_cuda_engine,TRTbin)
    def __create_context(self):
        for binding in self.__engine:
            size = trt.volume(self.__engine.get_binding_shape(binding)) * \
                   self.__engine.max_batch_size
            host_mem = cuda.pagelocked_empty(size, np.float32)
            cuda_mem = cuda.mem_alloc(host_mem.nbytes)
            self.__bindings.append(int(cuda_mem))
            if self.__engine.binding_is_input(binding):
                self.__host_inputs.append(host_mem)
                self.__cuda_inputs.append(cuda_mem)
            else:
                self.__host_outputs.append(host_mem)
                self.__cuda_outputs.append(cuda_mem)
        return self.__engine.create_execution_context()
    def __initcuda(self):
        self.__trt_logger = trt.Logger(trt.Logger.ERROR)
        self.__cuda_ctx = cuda.Device(0).make_context()
        self.__load_plugins()
        self.__engine = self.__load_engine()
        self.__stream = cuda.Stream()
        self.__host_inputs = []
        self.__cuda_inputs = []
        self.__host_outputs = []
        self.__cuda_outputs = []
        self.__bindings = []
        self.__context = self.__create_context()

    def __init__(self, model,input_wh=(300,300),videoCapture=None):
        """Initialize TensorRT plugins, runtime, engine and conetxt."""
        Thread.__init__(self)
        self.model=model
        self.cap=videoCapture
        self.__locker=Lock()
        self.__not_empty=Condition(self.__locker)
        self.__not_full=Condition(self.__locker)
        self.__runing=True
        self.__results=0
        self.__orgimg=None
        self.start()

    def run(self):
        self.__initcuda()
        self.cap.open()
        while self.__runing:
            img,orgimg = self.cap.read()
            cuda.memcpy_htod_async( self.__cuda_inputs[0],img.ravel(), self.__stream)
            self.__context.execute_async(  batch_size=1, bindings=self.__bindings,  stream_handle=self.__stream.handle)
            cuda.memcpy_dtoh_async(self.__host_outputs[1], self.__cuda_outputs[1], self.__stream)
            cuda.memcpy_dtoh_async(self.__host_outputs[0], self.__cuda_outputs[0], self.__stream)
            self.__stream.synchronize()
            with self.__not_full:
                while self.__results:
                    self.__not_full.wait()
                self.__orgimg=orgimg
                self.__results=1
                self.__not_empty.notify()

        del self.__stream
        del self.__cuda_inputs
        del self.__cuda_outputs
        del self.__host_inputs
        del self.__host_outputs
        self.__cuda_ctx.pop()
        del self.__cuda_ctx
    
    def stop(self):
        self.__runing=False
        with self.__not_empty:
            self.__results=0
            self.__not_full.notify()
        self.join()


    def getresult(self,conf_th=0.3):
        boxes, confs, clss = [], [], []
        with self.__not_empty:
            while not self.__results:
                self.__not_empty.wait()
            self.__results=0
            self.__not_full.notify()
        img=self.__orgimg  
        output=self.__host_outputs[0][:]
       
        for prefix in range(0, len(output),7):
            conf = float(output[prefix+2])
            if conf < conf_th:
                continue
            x1 = int(output[prefix+3] * 640)
            y1 = int(output[prefix+4] * 640)-80
            x2 = int(output[prefix+5] * 640)
            y2 = int(output[prefix+6] * 640)-80
            cls = int(output[prefix+1])
            a,b=x2-x1,y2-y1
            if (a*b)<2500: continue
            boxes.append((x1, y1, x2, y2))
            confs.append(conf)
            clss.append(cls)
        return img,boxes,confs,clss
